{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "\n",
    "from pymilvus import (\n",
    "    MilvusClient,\n",
    "    DataType,\n",
    "    Function,\n",
    "    FunctionType,\n",
    "    AnnSearchRequest,\n",
    "    RRFRanker,\n",
    ")\n",
    "\n",
    "from pymilvus.model.hybrid import BGEM3EmbeddingFunction\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 定义搜索器\n",
    "HybridRetriever 包括了四个部分。\n",
    "- __init__：定义一个HybridRetriever需要传的基本参数，例如uri, collection_name，稠密embedding时调用的function。\n",
    "- build_collection：包括以下三个部分。\n",
    "    - tokenizer_params：设置token化的参数（例如停用词），\n",
    "    - MilvusClient.create_schema()：创建schema和add_field，\n",
    "    - 还有\n",
    "- insert_data：把embedding之后的dense vector插入到milvus指定collection之中。\n",
    "- search：定义sparse、dense、hybrid等三种search方式。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class HybridRetriever:\n",
    "    def __init__(self, uri, collection_name=\"hybrid\", dense_embedding_function=None):\n",
    "        self.uri = uri\n",
    "        self.collection_name = collection_name\n",
    "        self.embedding_function = dense_embedding_function\n",
    "        self.use_reranker = True\n",
    "        self.use_sparse = True\n",
    "        self.client = MilvusClient(uri=uri)\n",
    "\n",
    "    def build_collection(self):\n",
    "        if isinstance(self.embedding_function.dim, dict):\n",
    "            dense_dim = self.embedding_function.dim[\"dense\"]\n",
    "        else:\n",
    "            dense_dim = self.embedding_function.dim\n",
    "\n",
    "        tokenizer_params = {\n",
    "            \"tokenizer\": \"standard\",\n",
    "            \"filter\": [\n",
    "                \"lowercase\",\n",
    "                {\n",
    "                    \"type\": \"length\",\n",
    "                    \"max\": 200,\n",
    "                },\n",
    "                {\"type\": \"stemmer\", \"language\": \"english\"},\n",
    "                {\n",
    "                    \"type\": \"stop\",\n",
    "                    \"stop_words\": [\n",
    "                        \"a\",\n",
    "                        \"an\",\n",
    "                        \"and\",\n",
    "                        \"are\",\n",
    "                        \"as\",\n",
    "                        \"at\",\n",
    "                        \"be\",\n",
    "                    ],\n",
    "                },\n",
    "            ],\n",
    "        }\n",
    "\n",
    "        schema = MilvusClient.create_schema()\n",
    "        schema.add_field(\n",
    "            field_name=\"pk\",\n",
    "            datatype=DataType.VARCHAR,\n",
    "            is_primary=True,\n",
    "            auto_id=True,\n",
    "            max_length=100,\n",
    "        )\n",
    "        schema.add_field(\n",
    "            field_name=\"content\",\n",
    "            datatype=DataType.VARCHAR,\n",
    "            max_length=65535,\n",
    "            analyzer_params=tokenizer_params,\n",
    "            enable_match=True,\n",
    "            enable_analyzer=True,\n",
    "        )\n",
    "        schema.add_field(\n",
    "            field_name=\"sparse_vector\", datatype=DataType.SPARSE_FLOAT_VECTOR\n",
    "        )\n",
    "        schema.add_field(\n",
    "            field_name=\"dense_vector\", datatype=DataType.FLOAT_VECTOR, dim=dense_dim\n",
    "        )\n",
    "        schema.add_field(\n",
    "            field_name=\"original_uuid\", datatype=DataType.VARCHAR, max_length=128\n",
    "        )\n",
    "        schema.add_field(field_name=\"doc_id\", datatype=DataType.VARCHAR, max_length=64)\n",
    "        schema.add_field(\n",
    "            field_name=\"chunk_id\", datatype=DataType.VARCHAR, max_length=64\n",
    "        ),\n",
    "        schema.add_field(field_name=\"original_index\", datatype=DataType.INT32)\n",
    "\n",
    "        functions = Function(\n",
    "            name=\"bm25\",\n",
    "            function_type=FunctionType.BM25,\n",
    "            input_field_names=[\"content\"],\n",
    "            output_field_names=\"sparse_vector\",\n",
    "        )\n",
    "\n",
    "        schema.add_function(functions)\n",
    "\n",
    "        index_params = MilvusClient.prepare_index_params()\n",
    "        index_params.add_index(\n",
    "            field_name=\"sparse_vector\",\n",
    "            index_type=\"SPARSE_INVERTED_INDEX\",\n",
    "            metric_type=\"BM25\",\n",
    "        )\n",
    "        index_params.add_index(\n",
    "            field_name=\"dense_vector\", index_type=\"FLAT\", metric_type=\"IP\"\n",
    "        )\n",
    "\n",
    "        self.client.create_collection(\n",
    "            collection_name=self.collection_name,\n",
    "            schema=schema,\n",
    "            index_params=index_params,\n",
    "        )\n",
    "\n",
    "    def insert_data(self, chunk, metadata):\n",
    "        embedding = self.embedding_function([chunk])\n",
    "        if isinstance(embedding, dict) and \"dense\" in embedding:\n",
    "            dense_vec = embedding[\"dense\"][0]\n",
    "        else:\n",
    "            dense_vec = embedding[0]\n",
    "        self.client.insert(\n",
    "            self.collection_name, {\"dense_vector\": dense_vec, **metadata}\n",
    "        )\n",
    "\n",
    "    def search(self, query: str, k: int = 20, mode=\"hybrid\"):\n",
    "\n",
    "        output_fields = [\n",
    "            \"content\",\n",
    "            \"original_uuid\",\n",
    "            \"doc_id\",\n",
    "            \"chunk_id\",\n",
    "            \"original_index\",\n",
    "        ]\n",
    "        if mode in [\"dense\", \"hybrid\"]:\n",
    "            embedding = self.embedding_function([query])\n",
    "            if isinstance(embedding, dict) and \"dense\" in embedding:\n",
    "                dense_vec = embedding[\"dense\"][0]\n",
    "            else:\n",
    "                dense_vec = embedding[0]\n",
    "\n",
    "        if mode == \"sparse\":\n",
    "            results = self.client.search(\n",
    "                collection_name=self.collection_name,\n",
    "                data=[query],\n",
    "                anns_field=\"sparse_vector\",\n",
    "                limit=k,\n",
    "                output_fields=output_fields,\n",
    "            )\n",
    "        elif mode == \"dense\":\n",
    "            results = self.client.search(\n",
    "                collection_name=self.collection_name,\n",
    "                data=[dense_vec],\n",
    "                anns_field=\"dense_vector\",\n",
    "                limit=k,\n",
    "                output_fields=output_fields,\n",
    "            )\n",
    "        elif mode == \"hybrid\":\n",
    "            full_text_search_params = {\"metric_type\": \"BM25\"}\n",
    "            full_text_search_req = AnnSearchRequest(\n",
    "                [query], \"sparse_vector\", full_text_search_params, limit=k\n",
    "            )\n",
    "\n",
    "            dense_search_params = {\"metric_type\": \"IP\"}\n",
    "            dense_req = AnnSearchRequest(\n",
    "                [dense_vec], \"dense_vector\", dense_search_params, limit=k\n",
    "            )\n",
    "\n",
    "            results = self.client.hybrid_search(\n",
    "                self.collection_name,\n",
    "                [full_text_search_req, dense_req],\n",
    "                ranker=RRFRanker(),\n",
    "                limit=k,\n",
    "                output_fields=output_fields,\n",
    "            )\n",
    "        else:\n",
    "            raise ValueError(\"Invalid mode\")\n",
    "        return [\n",
    "            {\n",
    "                \"doc_id\": doc[\"entity\"][\"doc_id\"],\n",
    "                \"chunk_id\": doc[\"entity\"][\"chunk_id\"],\n",
    "                \"content\": doc[\"entity\"][\"content\"],\n",
    "                \"score\": doc[\"distance\"],\n",
    "            }\n",
    "            for doc in results[0]\n",
    "        ]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 插入数据\n",
    "\n",
    "这里使用codebase_chunks.json的示例数据作为演示。\n",
    "\n",
    "有若干个docs，每个doc有若干个chunks。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "path = \"codebase_chunks.json\"\n",
    "with open(path, \"r\") as f:\n",
    "    dataset = json.load(f)\n",
    "\n",
    "is_insert = True\n",
    "if is_insert:\n",
    "    standard_retriever.build_collection()\n",
    "    for doc in dataset:\n",
    "        doc_content = doc[\"content\"]\n",
    "        for chunk in doc[\"chunks\"]:\n",
    "            metadata = {\n",
    "                \"doc_id\": doc[\"doc_id\"],\n",
    "                \"original_uuid\": doc[\"original_uuid\"],\n",
    "                \"chunk_id\": chunk[\"chunk_id\"],\n",
    "                \"original_index\": chunk[\"original_index\"],\n",
    "                \"content\": chunk[\"content\"],\n",
    "            }\n",
    "            chunk_content = chunk[\"content\"]\n",
    "            # standard_retriever.insert_data(chunk_content, metadata)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
